Deliverable 4: the slimmed platform-optimizer convention (contracts, portable state, DCP) - #67

OptimizerCapabilities accepted arbitrary objects in training, checkpoints,
and precisions and non-bool capability flags, and OptimizerChildContract
accepted any object as its child contract, so an invalid declaration
constructed silently and failed far from where it was built. Check element
and flag types in __post_init__ like every sibling descriptor.

_hybrid_contract copied the muon child's training claims verbatim and fell
back to the static base tuple when muon was absent, never consulting the
backup child. A Gefen-backed backup-only hybrid therefore omitted the
flattened element-shard layout its backup validates, while an AdamW-backed
hybrid claimed DTensor training no code validates for that child.

Declare the ordered union of the present children's own claims instead;
a backup without a contract contributes only plain replicated training.

The canonical-import and hybrid-rebinding atomicity tests asserted only
top-level object identity, so a regression that published staged state by
mutating live containers in place (state[param].update, copy_() into an
existing state tensor, group option or group['params'] element edits)
would pass every assertion. Add tests/_state_snapshot.py, a shared deep
snapshot helper that captures container identities plus bitwise clones of
every reachable tensor, per-parameter state dict, counter, and param-group
entry, and rewire both files' snapshot helpers onto it. Verified that all
eight simulated in-place regressions now fail the assertions while the
current implementation still passes.

test_activation_and_restore_copy_failures_are_atomic only checked that
optimizer.state, the per-parameter dict, and the codebook were the same
objects after an injected copy failure; an offload/restore regression that
partially overwrote m_codebook/m_magnitude/vmean in place before failing
would keep every identity intact and pass. Snapshot the persistent tensors
and counters with the file's existing _persistent_snapshot helper before
each injected failure and assert bitwise equality (plus codebook value and
global step) afterwards. Verified the new assertions catch a simulated
in-place corruption that the old identity checks missed.

Every flattened shard in the scoped-collective tests had nonzero length,
leaving the dedicated empty-shard logic (nonempty_activity filtering of
the gradient-presence consensus and the empty-slice no-op paths) entirely
unexercised. Add a two-member gloo test where one member binds a legal
zero-length flattened slice to a live numel-0 tensor and assert it joins
initialization, step, refresh, and failure-sync collectives symmetrically:
codebooks agree with the nonempty-only oracle on both members, presence
asymmetry on the empty member is accepted, the empty member's state stays
inert, a nonempty-member exact-DP failure is synchronized to the empty
member and leaves both atomic and retryable.

The reuse_existing_periods flag that gates the conditional "initialize"
operation header in _prepare_gefen_exact_codebook was derived from the
rank-local _resuming_from_checkpoint() predicate on both hot paths
(_maybe_refresh_gefen_codebook and initialize_codebook). Per-parameter
state presence is legitimately rank-asymmetric under an explicit scope
(empty non-owner slots and zero-length flattened shards never create
state), so after a consolidation-style resume that strips the common
codebook, state-bearing members computed reuse=True and skipped the
header all_gather while empty members computed reuse=False and issued
it, mismatching the group's collective schedules.

Resolve one group-wide decision before any member branches on it:
_scope_agreed_resuming_from_checkpoint all_reduces (MAX) the local
predicate over the codebook scope, so any member that restored periods
makes the whole scope reuse them. Both call sites are reached by every
member of a multi-member scope in the same order, and empty members are
unaffected behaviorally by reuse=True because parameters without
gradients or elements never enter the period iterator.

_ensure_codebook_scope_agreement early-returned on the rank-local
_gefen_codebook_scope_validated flag. Collective-free rank-local
operations (move_state_, offload_state_, staged and native checkpoint
loads) reset that flag on the member that ran them only, while every
value the step operation header fingerprints stays identical (the
codebook fingerprint is computed on CPU bytes and is device
independent). One member then entered the agreement all_gather while
the others early-returned, diverging the scoped collective schedule
directly behind a passing header exchange.

Fold the flag into the always-exchanged operation header as a trailing
decision bit that is excluded from the equality check, and lower the
flag on every member when any member reports it cleared, so the group
re-validates together and then proceeds with identical collectives.

This resolution was chosen over the two alternatives deliberately:
raising on a flag mismatch would turn the advertised collective-free
rank-local operations into scope-wide faults on legitimately
asymmetric use, and dropping the reset for movement/offload would keep
the divergence for staged and native loads, which genuinely can change
scope-agreement-relevant state and must keep resetting the flag.

_canonical_import_live_token identified locally bound parameters by
id() only, so commit_canonical_state_import could not detect that a
parameter's storage was retargeted (or mutated in place) between
prepare and commit, even though the staged shadow was device-cast and
geometry-validated against the parameters as they existed at prepare
time. The sibling portable path already folds _parameter_storage_token
(device, dtype, layout, shape, stride, storage pointer/size, version)
into its live token for every local binding; mirror it here so a stale
prepared import is refused with the existing freshness error.

Cover the three review findings: a consolidation-style resume with a
whole-parameter owner and an empty non-owner must gate the conditional
initialize header on one group-wide reuse decision and then complete
the resume step collectively with restored periods intact; a
collective-free move_state_ on one member only must lead every member
to the same scope re-validation decision on the next step; and a
prepared canonical import must go stale when a bound parameter's
storage is retargeted or mutated between prepare and commit, while an
undisturbed prepare/commit round still succeeds.

…k scope

GefenMuonHybrid.step decided the GradScaler overflow skip with the
rank-local _amp_prepare_optimizer_step and raised structural gradient
preflight errors rank-locally, even when post_sharding installed one
multi-member codebook process-group binding shared by both children. A
rank whose found_inf was set (or whose local gradients were malformed)
then skipped or raised alone while its peers entered the children's
scoped step collectives, hanging the group.

Route both decisions through the children's scoped protocol before any
child steps: the composite runs Gefen._prepare_scoped_amp_optimizer_step
on itself (found_inf/grad_scale agreement is validated collectively and
an overflow skip is a group-wide decision, entered and exited
symmetrically on every member) and synchronizes the structural preflight
through the children's _synchronize_codebook_scope_failure so every
scope member raises together, keeping the atomic both-children-skip
semantics. Without a binding the local behavior is unchanged.

step() re-ran the complete O(params x world) finalized-layout forensic
rebuild on every guard call (2 passes per unscoped step, up to 7 under a
multi-member codebook scope) and recomputed the manifest sha256
fingerprint inside every scoped operation header, costing seconds of
host time per step at large scale.

Cache one forensic verdict as an O(local params) identity-token snapshot
(finalized registries by object identity, every live group container,
parameter and compatibility name, plus a version counter bumped by every
legitimate mutating API), and compute the manifest shard set and digest
once per finalized manifest at post_sharding. Steady-state step guards
now reuse the verdict; checkpoint prepare/commit, canonical
export/import, rebinding, state movement/offload, codebook
initialize/refresh, scope re-validation, and contract readiness still
run the full forensic rebuild. The offload step-readiness scan is
deduped under the same scheme. Caches live in __slots__ so staged
__dict__ copies and fail-before-mutation snapshots never see them.

Public-container tampering (group params/param_names slots, state
names, the name cache), including closure-time mutation, still fails
the step guard before any state mutation; in-place edits inside the
private registries move to detection at the next full-forensics
boundary, as now documented in docs/optimizer_contracts.md.

Under active state offload the per-step readiness guard
(_assert_state_offload_step_ready, run ~2x/step) reached
_state_movement_rejection_reason, which always consulted the finalized
layout with full=True. That forced an uncached full layout rebuild plus a
manifest sha256 recompute on every step (two full passes and two digest
recomputes per step), bypassing the per-step layout-forensics memoization
the non-offload guards already use.

The finalized binding layout is immutable across steps, so thread a
require_full_layout flag through _state_offload_rejection_reason and
_state_movement_rejection_reason. The per-step offload readiness path uses
the memoized fast path (full=False, reusing _gefen_layout_forensics_verdict),
while the move_state_ / _atomic_state_movement_supported boundaries,
activation, load, and contract readiness keep the full rebuild. The
per-tensor offload state scan (CPU tightness/dtype, pairwise disjointness,
native-schema validation) is untouched and still runs on every step, so a
token-preserving corruption of a later parameter's offloaded state is still
caught before any earlier parameter mutates.

The fail-before-mutation snapshot compared raw bytes via view(torch.uint8),
which PyTorch rejects on a 0-dim tensor ("self.dim() cannot be 0 to view
Float as Byte"). Capturable device-resident scalar counters hit this, failing
the CUDA-only capturable canonical-import tests (which the CPU CI skips).
Flatten to 1-D before reinterpreting; shapes are already checked equal, so
both sides flatten identically.

…ed load atomicity)

main now carries the fail-before-mutation load_state_dict work (Tier-1, #69), so
it no longer needs to live in this convention branch. Merging main in subtracts
that shared base from this PR's diff: the two load-atomicity test files become
identical to main, and the native/hybrid staging split shows up only as this
branch's convention-only decorations rather than as net-new code.

Reconcile the atomicity methods, which had diverged across the two branches:

  * main has the review-fixed staging -- isolated foreign-AdamW backup staging
    (_stage_foreign_backup_load / _commit_foreign_backup_load) because torch
    AdamW.load_state_dict is not itself fail-before-mutation, plus the
    _run_load_state_dict_post_hooks split so a composite owner commits every
    child's raw state before dispatching any child's post-hook -- but none of
    the convention decorations;
  * this branch has the decorations (the _assert_finalized_binding_layout
    guards, the _stage_load_state_dict offload block, the
    _invalidate_layout_forensics_caches commit hook, and the
    allow_preinitialized_periods canonical validation) layered on the older,
    pre-review hybrid load.

Keep both: this branch's decorations on top of main's review-fixed base. The
foreign-backup helpers and the three-phase hybrid load body are byte-identical
to main; gefen.py differs from this branch only by the post-hook-split addition.

tests/test_native_load_atomicity.py + tests/test_hybrid_load_atomicity.py:
28 passed against the reconciled tree.